import os
import time
from collections import deque
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import uvicorn
from datetime import datetime

app = FastAPI(title="Directory Traversal API", 
              description="高效遍历Windows系统目录的API服务", 
              version="1.0.0")


# 定义标准响应模型
class ResponseModel(BaseModel):
    """
    标准响应模型
    
    Attributes:
        msg: 响应消息
        code: 状态码
        data: 响应数据
    """
    msg: str
    code: int
    data: Optional[Any] = None

class FileInfo(BaseModel):
    path: str
    name: str
    size: int
    modified_time: str
    created_time: str
    extension: Optional[str] = None

class TraversalResponse(BaseModel):
    count: int
    total_size: int
    traversal_time: float
    files: List[FileInfo]
    directory: str

def format_timestamp(timestamp: float) -> str:
    """将时间戳转换为自定义格式的字符串（没有T字符）"""
    dt = datetime.fromtimestamp(timestamp)
    # 格式化为: YYYY-MM-DD HH:MM:SS.microseconds
    return dt.strftime("%Y-%m-%d %H:%M:%S.%f")

def is_after_target_date(modified_time_str: str, target_date: str = "2025-09-12") -> bool:
    """检查修改时间是否在目标日期之后"""
    # 提取日期部分进行比较
    file_date = modified_time_str.split(" ")[0]  # 获取YYYY-MM-DD部分
    return file_date > target_date

def traverse_files(root_dir, filter_date: str = None):
    """
    使用广度优先搜索（BFS）非递归遍历目录中的所有文件（不含目录）
    返回生成器以节省内存并提高效率
    
    filter_date: 筛选条件，只返回修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    queue = deque([root_dir])
    while queue:
        current_dir = queue.popleft()
        try:
            with os.scandir(current_dir) as entries:
                for entry in entries:
                    try:
                        if entry.is_dir(follow_symlinks=False):
                            # 如果是目录，添加到队列中以便遍历其子内容
                            queue.append(entry.path)
                            continue  # 跳过目录，不将其作为文件返回
                        
                        # 处理文件
                        stat = entry.stat()
                        modified_time_str = format_timestamp(stat.st_mtime)
                        
                        # 如果有筛选日期，检查是否满足条件
                        if filter_date and not is_after_target_date(modified_time_str, filter_date):
                            # 不满足筛选条件，跳过此文件
                            continue
                        
                        file_info = {
                            "path": entry.path,
                            "name": entry.name,
                            "size": stat.st_size,
                            "modified_time": modified_time_str,
                            "created_time": format_timestamp(stat.st_ctime),
                            "extension": os.path.splitext(entry.name)[1] if '.' in entry.name else None
                        }
                        
                        yield file_info
                    except (OSError, PermissionError):
                        # 跳过无法访问的文件或目录
                        continue
        except PermissionError:
            # 跳过权限不足的目录
            continue
        except OSError as e:
            # 处理其他可能的OS错误
            print(f"Error accessing {current_dir}: {e}")
            continue

@app.get("/traverse", response_model=ResponseModel, 
         summary="遍历指定目录", 
         description="递归遍历指定目录下的所有文件（不含目录）")
async def traverse_directory(path: Optional[str] = None, modified_after: Optional[str] = None):
    """
    遍历指定目录并返回所有文件信息（不含目录）
    
    - **path**: 要遍历的目录路径，默认为当前工作目录
    - **modified_after**: 只返回修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    # 验证日期格式
    if modified_after:
        try:
            datetime.strptime(modified_after, "%Y-%m-%d")
        except ValueError:
            raise HTTPException(status_code=400, detail="日期格式不正确，请使用YYYY-MM-DD格式")
    
    start_time = time.time()
    files = []
    total_size = 0
    
    # 遍历目录并收集文件信息
    for file_info in traverse_files(path, modified_after):
        files.append(file_info)
        total_size += file_info["size"]
    
    traversal_time = time.time() - start_time
    

    return ResponseModel(
            msg="ok",
            code=200,
            data={
                "count": len(files),
                "total_size": total_size,
                "traversal_time": traversal_time,
                "files": files,
                "directory": path}
        )


@app.get("/stats", summary="获取目录统计信息")
async def get_directory_stats(path: Optional[str] = None, modified_after: Optional[str] = None):
    """
    获取目录的统计信息而不返回所有文件列表
    
    - **path**: 要统计的目录路径，默认为当前工作目录
    - **modified_after**: 只统计修改时间在此日期之后的文件（格式：YYYY-MM-DD）
    """
    if path is None:
        path = os.getcwd()
    
    # 验证路径是否存在
    if not os.path.exists(path):
        raise HTTPException(status_code=404, detail="指定路径不存在")
    
    # 验证路径是否为目录
    if not os.path.isdir(path):
        raise HTTPException(status_code=400, detail="指定路径不是目录")
    
    # 验证日期格式
    if modified_after:
        try:
            datetime.strptime(modified_after, "%Y-%m-%d")
        except ValueError:
            raise HTTPException(status_code=400, detail="日期格式不正确，请使用YYYY-MM-DD格式")
    
    start_time = time.time()
    file_count = 0
    total_size = 0
    
    # 使用os.walk进行更高效的统计（不收集所有文件信息）
    for root, dirs, files_list in os.walk(path):
        # 筛选文件
        if modified_after:
            for file in files_list:
                file_path = os.path.join(root, file)
                try:
                    mtime = os.path.getmtime(file_path)
                    if datetime.fromtimestamp(mtime).strftime("%Y-%m-%d") > modified_after:
                        file_count += 1
                        total_size += os.path.getsize(file_path)
                except (OSError, PermissionError):
                    continue
        else:
            file_count += len(files_list)
            for file in files_list:
                try:
                    file_path = os.path.join(root, file)
                    total_size += os.path.getsize(file_path)
                except (OSError, PermissionError):
                    continue
    
    traversal_time = time.time() - start_time
    
    return {
        "directory": path,
        "file_count": file_count,
        "total_size": total_size,
        "traversal_time": traversal_time,
        "filter_date": modified_after
    }

@app.get("/", summary="API信息")
async def root():
    return {
        "message": "目录遍历API服务",
        "version": "1.0.0",
        "endpoints": {
            "/traverse": "遍历指定目录下的所有文件（不含目录，返回详细信息）",
            "/stats": "获取目录统计信息（不返回文件列表）",
            "/docs": "API文档"
        }
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)